package com.small.xx.ai.spring.rag.controller;

import com.small.xx.ai.spring.rag.config.Constant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.DefaultContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.ai.model.transformer.SummaryMetadataEnricher;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.transformer.ContentFormatTransformer;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RestController
@RequestMapping("/document/transformer")
public class DocumentTransformerController {

    private static final Logger logger = LoggerFactory.getLogger(DocumentTransformerController.class);

    private final List<Document> documents;
    private final ChatModel chatModel;

    /**
     * 构造方法：初始化时读取PDF文件并加载文档内容
     *
     * @param chatModel 注入的聊天模型实例，用于后续AI增强操作
     */
    public DocumentTransformerController(ChatModel chatModel) {
        // 日志记录：开始按页读取PDF文件
        logger.info("start read pdf file by page");
        // 通过资源加载器获取PDF文件资源（路径由Constant.PDF_FILE_PATH指定）
        Resource resource = new DefaultResourceLoader().getResource(Constant.PDF_FILE_PATH);
        // 创建按页读取PDF的阅读器（仅支持PDF格式文件）
        PagePdfDocumentReader pagePdfDocumentReader = new PagePdfDocumentReader(resource);
        // 读取PDF内容，按页拆分为Document对象列表并存储
        this.documents = pagePdfDocumentReader.read();

        // 初始化聊天模型实例
        this.chatModel = chatModel;
    }

    /**
     * 基于Token的文本分块接口：将文档按Token数量拆分，适配大模型上下文窗口
     *
     * @return 拆分后的文档块列表
     */
    @GetMapping("/token-text-splitter")
    public List<Document> tokenTextSplitter() {
        // 日志记录：开始执行Token文本分块
        logger.info("start token text splitter");
        // 构建Token文本分块器，配置分块参数
        TokenTextSplitter tokenTextSplitter =
            TokenTextSplitter.builder().withChunkSize(800)                 // 每个文本块的目标Token数量（如800 Token）
                .withMinChunkSizeChars(350)         // 每个文本块的最小字符数（避免过短块）
                .withMinChunkLengthToEmbed(5)       // 丢弃字符数小于此值的文本块（过滤无效内容）
                .withMaxNumChunks(10000)            // 生成的最大文本块数量（防止过度拆分）
                .withKeepSeparator(true)            // 保留文本中的分隔符（如段落标记）
                .build();
        // 对原始文档执行分块操作并返回结果
        return tokenTextSplitter.split(this.documents);
    }

    /**
     * 内容格式转换接口：统一文档内容的格式（如去除多余空格、标准化换行等）
     *
     * @return 格式标准化后的文档列表
     */
    @GetMapping("/content-format-transformer")
    public List<Document> contentFormatTransformer() {
        // 日志记录：开始执行内容格式转换
        logger.info("start content format transformer");
        // 创建默认内容格式化器（使用默认配置，如处理空白字符、统一换行符等）
        DefaultContentFormatter defaultContentFormatter = DefaultContentFormatter.defaultConfig();

        // 创建内容格式转换器，传入格式化器
        ContentFormatTransformer contentFormatTransformer = new ContentFormatTransformer(defaultContentFormatter);

        // 对原始文档执行格式转换并返回结果
        return contentFormatTransformer.apply(this.documents);
    }

    /**
     * 关键词元数据增强接口：为文档添加关键词元数据（通过AI模型提取）
     *
     * @return 包含关键词元数据的文档列表
     */
    @GetMapping("/keyword-metadata-enricher")
    public List<Document> keywordMetadataEnricher() {
        // 日志记录：开始执行关键词元数据增强
        logger.info("start keyword metadata enricher");
        // 创建关键词元数据增强器：使用聊天模型，为每个文档提取3个关键词
        KeywordMetadataEnricher keywordMetadataEnricher = new KeywordMetadataEnricher(this.chatModel, 3);
        // 为原始文档添加关键词元数据并返回结果
        return keywordMetadataEnricher.apply(this.documents);
    }

    /**
     * 摘要元数据增强接口：为文档添加摘要元数据（包括当前、前序、后续文档的摘要）
     *
     * @return 包含摘要元数据的文档列表
     */
    @GetMapping("/summary-metadata-enricher")
    public List<Document> summaryMetadataEnricher() {
        // 日志记录：开始执行摘要元数据增强
        logger.info("start summary metadata enricher");
        // 定义需要生成的摘要类型：当前文档摘要、前序文档摘要、后续文档摘要
        List<SummaryMetadataEnricher.SummaryType> summaryTypes =
            List.of(SummaryMetadataEnricher.SummaryType.NEXT,    // 后续文档摘要
                SummaryMetadataEnricher.SummaryType.CURRENT, // 当前文档摘要
                SummaryMetadataEnricher.SummaryType.PREVIOUS);// 前序文档摘要
        // 创建摘要元数据增强器：使用聊天模型，按指定类型生成摘要
        SummaryMetadataEnricher summaryMetadataEnricher = new SummaryMetadataEnricher(this.chatModel, summaryTypes);

        // 为原始文档添加摘要元数据并返回结果
        return summaryMetadataEnricher.apply(this.documents);
    }
}
